from Crunchers import kdtree
from Einput import einput
from numpy import array, ndarray

class Neighbors(object):
    """ Wrapper around the Kd-tree cython module to work with entities.
    """
    def __init__(self, entities, level ='A', ctype =None, bucket_size =10):
        self.entities = einput(entities, level, name =level)
        self.ctype  = (ctype or 'coords')                         # coords type.
        self.coords = array(self.entities._data_children(self.ctype))
        self.kdt = kdtree(self.coords, bucket_size)               # kdtree from array
        self.indices = dict(enumerate(self.entities.keys()))      # indices -> id's

    def get_neighbors(self, entity, r_or_k, level ='A'):
        if not isinstance(entity, ndarray):                       # grab query coordinate
            entity = getattr(entity, self.ctype)
        if not isinstance(r_or_k, int):
            (indices, distances) = self.kdt.rn(entity, r_or_k*r_or_k)  # find within r
        else:
            (indices, distances) = self.kdt.knn(entity, r_or_k)        # find k-neighbors
        entities = []
        for i in indices:                                         # only one query
            entities.append(self.entities[self.indices[i]])       # indice -> id's -> entity
        return einput(entities, level, name ='neighbours')



"""
class Neighbors(object):

    def __init__(self, entities, level ='A', coords ='coords', bucket_size =10):
        self.entities = einput(entities, level, name = level)
        self.ctype = coords # coords type
        self.coords = array(self.entities._data_children(coords))
        self.indices = dict(enumerate(self.entities.keys()))    # indices -> id's
        self.kdt = KDTree(self.coords.shape[1], bucket_size)
        self.kdt.set_coords(self.coords)

    def get_neighbors(self, entity, radius, level ='A'):
        if not isinstance(entity, ndarray):
            entity = getattr(entity, self.ctype)
        self.kdt.search(entity, radius)
        indices = self.kdt.get_indices()
        entities = []
        for i in indices:
            entities.append(self.entities[self.indices[i]])     # indice -> id's -> entity
        return einput(entities, level, name ='neighbors')

    def get_pairs(self, radius, level ='A'):
        self.kdt.all_search(radius)
        indices = self.kdt.all_get_indices()

        entity_pairs = set()
        for i1, i2 in indices:
            a1 = self.entities[self.indices[i1]].get_parent(level)
            a2 = self.entities[self.indices[i2]].get_parent(level)
            pair = (a1, a2)
            entity_pairs.add(pair)

        return entity_pairs


class KNNeighbors(object):
    def __init__(self, entities, level ='A', ctype ='coords'):
        self.entities = einput(entities, level, name =level)
        self.ctype  = ctype                                       # coords type.
        self.coords = array(self.entities._data_children(self.ctype))
        self.kdt = ann.kdtree(self.coords)                        # kdtree from array / ANN
        self.indices = dict(enumerate(self.entities.keys()))      # indices -> id's

    def get_neighbors(self, entity, k, level ='A'):
        if not isinstance(entity, ndarray):
            entity = getattr(entity, self.ctype)
        (indices, distances) = self.kdt.knn(entity, k)  # find k
        entities = []
        for i in indices[0]: # only one query
            entities.append(self.entities[self.indices[i]])          # indice -> id's -> entity
        return einput(entities, level, name ='neighbors')
"""